import os
import itertools
from typing import Iterator, Optional, Generator

import streamlit as st
from dotenv import load_dotenv

# 通过.env文件设置环境变量
# reference: https://github.com/theskumar/python-dotenv
load_dotenv()

import api
from data_types import TextMsg, ImageMsg, TextMsgList, MsgList, filter_text_msg

debug = os.getenv("DEBUG", "").lower() in ("1", "yes", "y", "true", "t", "on")


st.set_page_config(page_title="Role Paly Demo", page_icon="🤖", layout="wide")


def update_api_key(key: Optional[str] = None):
    if debug:
        print(
            f'update_api_key. st.session_state["API_KEY"] = {st.session_state["API_KEY"]}, key = {key}'
        )
    key = key or st.session_state["API_KEY"]
    if key:
        api.API_KEY = key


# 设置API KEY
api_key = st.sidebar.text_input(
    "API_KEY",
    value=os.getenv("API_KEY", ""),
    key="API_KEY",
    type="password",
    on_change=update_api_key,
)
update_api_key(api_key)


if "role1_name" not in st.session_state:
    st.session_state.role1_name = ""

if "role1_characters" not in st.session_state:
    st.session_state.role1_characters = ""

if "role2_name" not in st.session_state:
    st.session_state.role2_name = ""

if "role2_characters" not in st.session_state:
    st.session_state.role2_characters = ""


def update_role1_name():
    st.session_state.role1_name = st.session_state.role1_name_val


# 初始化
if "history_role1" not in st.session_state:
    st.session_state["history_role1"] = []
    st.session_state["history_role2"] = []

if "meta" not in st.session_state:
    st.session_state["meta_role1"] = {
        "user_info": "",
        "bot_info": "",
        "bot_name": "",
        "user_name": "",
    }
    st.session_state["meta_role2"] = {
        "user_info": "",
        "bot_info": "",
        "bot_name": "",
        "user_name": "",
    }


def output_stream_response(response_stream: Iterator[str]):
    all_content = []
    content = ""
    for content in itertools.accumulate(response_stream):
        all_content.append(content)
    return content


def generate_character_profile(role_profile):
    return api.generate_role_appearance(role_profile)


def generate_role_name(role_profile):
    name = api.generate_name(role_profile).content
    if len(name) > 3:
        name = name[-3:]
    return name


def generate_dialogue(meta_role1, meta_role2, prologue, num: int):
    # 使用CharacterGLM交替生成两个角色的对话
    st.write("生成的对话:")
    # 追加开场白
    st.session_state["history_role1"].append(
        TextMsg({"role": "user", "content": prologue})
    )
    st.write(meta_role1["user_name"] + ":" + prologue)
    st.session_state["history_role2"].append(
        TextMsg({"role": "assistant", "content": prologue})
    )
    i = 1
    for _ in range(num):
        print("------------------第" + str(i) + "+轮----------------------------------------")
        response1 = api.get_characterglm_response(
            st.session_state["history_role1"], meta_role1
        )
        text1 = response_to_string(response1)
        st.session_state["history_role1"].append(
            TextMsg({"role": "assistant", "content": text1})
        )
        st.session_state["history_role2"].append(
            TextMsg({"role": "user", "content": text1})
        )
        st.write(meta_role1["user_name"] + ":" + text1)
        response2 = api.get_characterglm_response(
            st.session_state["history_role2"], meta_role2
        )
        text2 = response_to_string(response2)
        st.session_state["history_role1"].append(
            TextMsg({"role": "user", "content": text2})
        )
        st.session_state["history_role2"].append(
            TextMsg({"role": "assistant", "content": text2})
        )
        st.write(meta_role1["bot_name"] + ":" + text2)
        i = i + 1
        print("-----------------------------------------------------------")


def response_to_string(response):
    response_strings = []
    for line in response:
        response_strings.append(line)

    # 连接所有字符串
    dialog_string = "".join(response_strings)
    return dialog_string


@st.dialog("角色人设生成")
def open_dialog(roleNameAttr, roleCharactersAttr):
    text = st.text_area(
        "请描述角色特性:",
        max_chars=300,
        placeholder="可以是小说的一个段落，或者人物特性，比如样貌，穿着，性格等",
    )
    if st.button("生成人设和名称"):
        profileStream = generate_character_profile(text)
        roleName = generate_role_name(text)
        profileText = output_stream_response(profileStream)
        st.session_state[roleNameAttr] = roleName
        st.session_state[roleCharactersAttr] = profileText

        st.rerun()


def save_dialogue_to_file(dialogue, meta, filename="dialogue.txt", encoding="utf-8"):

    with open(filename, "w") as file:
        for message in dialogue:
            name = ""
            if message["role"] == "user":
                name = meta["user_name"]
            else:
                name = meta["bot_name"]
            text = message["content"]
            # 现在可以安全地将内容写入文件
            file.write(name + ":" + text + "\n")


# Streamlit网页布局
st.title("角色对话生成工具")

with st.container():
    col1, col2 = st.columns(2)
    # 文本框输入小说段落
    with col1:
        if st.button("生成角色1人设"):
            open_dialog("role1_name", "role1_characters")
        role1_name = st.text_input("角色1名称", key="role1_name")
        role1_characters = st.text_area(
            "角色1人设：", value="", key="role1_characters", height=100
        )

    with col2:
        if st.button("生成角色2人设"):
            open_dialog("role2_name", "role2_characters")
        role2_name = st.text_input("角色2名称", key="role2_name")
        role2_characters = st.text_area(
            "角色2人设：", key="role2_characters", height=100
        )


# 循环次数
number = st.number_input("生成几轮对话", value=5, min_value=1, max_value=10)
# 生成对话按钮
if st.button("生成对话"):
    meta = dict(
        user_name=st.session_state.role1_name,
        user_info=st.session_state.role1_characters,
        bot_name=st.session_state.role2_name,
        bot_info=st.session_state.role2_characters,
    )
    meta2 = dict(
        user_name=st.session_state.role2_name,
        user_info=st.session_state.role2_characters,
        bot_name=st.session_state.role1_name,
        bot_info=st.session_state.role1_characters,
    )
    prologue = api.generate_chat_prologue_prompt(meta)
    message = prologue.content
    if message.startswith(meta["user_name"] + ":"):
        message = message.split(":", 1)[1]

    # st.write("生成场景:")
    # st.write(message)

    if not isinstance(meta, dict):
        raise TypeError("meta must be a dictionary")

    generate_dialogue(meta, meta2, message, number)

    # st.write("生成的对话:")
    # for line in dialogue:
    #    st.write(line)

    # 保存对话到文件
    dialogue = st.session_state["history_role1"]
    save_dialogue_to_file(dialogue, meta)
    st.success("对话已保存到dialogue.txt文件中。")
